import torch
from model import ActionClassificationModel

if __name__ == '__main__':
    """
    input_shape: 模型输入shape
    ResNet50的参数应该与训练集保持一致
    pretrain：训练好的权重
    onnx_path：输出名字
    """
    # Define input shape
    input_shape = (1, 40, 3, 112, 112)

    # Load PyTorch model
    model = ActionClassificationModel()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    pretrain = 'pth/model_epoch_300.pth'
    model.load_state_dict(torch.load(pretrain, map_location=device))

    # Set model to evaluation mode
    model.eval()

    # Export PyTorch model to ONNX format
    dummy_input = torch.randn(*input_shape)
    onnx_path = 'pth/best_model.onnx'
    torch.onnx.export(model, dummy_input, onnx_path, verbose=True)
